from PIL import Image
import os

# 定义路径
normal_dir = '/root/code/fe-ma-sr/test/normal'
white_dir = '/root/code/fe-ma-sr/test/white'
output_dir = '/root/code/fe-ma-sr/test/cropped'

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 获取normal目录下的图片文件
normal_filenames = os.listdir(normal_dir)
print(f'Found {len(normal_filenames)} images in normal directory')

# 获取white目录下的所有图片文件名（用于命名）
white_filenames = os.listdir(white_dir)
print(f'Found {len(white_filenames)} images in white directory')

# 读取normal图片
normal_img_path = os.path.join(normal_dir, normal_filenames[0])
with Image.open(normal_img_path) as normal_img:
    print(f'Original normal image size: {normal_img.size}')
    
    # 为每个white目录中的图片创建裁剪版本
    for i, white_filename in enumerate(white_filenames):
        # 获取目标尺寸
        white_img_path = os.path.join(white_dir, white_filename)
        with Image.open(white_img_path) as white_img:
            target_width, target_height = white_img.size
            
            # 裁剪normal图片到目标尺寸
            # 使用居中裁剪
            cropped_img = normal_img.copy()
            if cropped_img.width > target_width and cropped_img.height > target_height:
                left = (cropped_img.width - target_width) // 2
                top = (cropped_img.height - target_height) // 2
                right = left + target_width
                bottom = top + target_height
                cropped_img = cropped_img.crop((left, top, right, bottom))
            elif cropped_img.width < target_width or cropped_img.height < target_height:
                # 如果normal图片小于目标尺寸，则调整大小后再裁剪
                # 先按比例放大到至少能覆盖目标尺寸
                scale = max(target_width / cropped_img.width, target_height / cropped_img.height)
                new_width = int(cropped_img.width * scale)
                new_height = int(cropped_img.height * scale)
                cropped_img = cropped_img.resize((new_width, new_height), Image.LANCZOS)
                
                # 再进行居中裁剪
                left = (cropped_img.width - target_width) // 2
                top = (cropped_img.height - target_height) // 2
                right = left + target_width
                bottom = top + target_height
                cropped_img = cropped_img.crop((left, top, right, bottom))
            
            # 保存裁剪后的图片，使用white目录中的文件名
            output_path = os.path.join(output_dir, white_filename)
            cropped_img.save(output_path)
            
            if i % 50 == 0:
                print(f'Processed {i+1}/{len(white_filenames)} images')

print('All images processed successfully!')